import os
import torch
from EAattak.Attck import Attck
from configure import configure
from data_utils import Dataer
from network import Network
import numpy as np
import time

from exper.ExpMUAPs import ExpMUAPs
from exper.ExpUntarget import ExpUntarget
from exper.ExpTarget import ExpTarget
from exper.ExpTransfer import ExpTransfer
from exper.ExpHighDim import ExpHighDim

from utils import Print_config, Plot
from FNSack.FNS import FNSack
from torchattacks import OnePixel
from SimpleAttack import SimBA
from torchattacks import FGSM, PGD

if __name__ == "__main__":

    config = vars(configure())

    ## os.environ['CUDA_VISBLE_DEVICES'] = '{}'.format(config['gpu_id'])
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    config['device'] = device

    Print_config(config=config)

    dataer = Dataer(dataset_name=config['dataset'], block_size=config['block_size'])

    neter = Network(dataer, config)
    atk = Attck(dataer=dataer, neter=neter, config=config)

    ## for FNS attack method
    ## test_loader = dataer.get_loader(batch_size=1, isTrain=False, isShuffle=True)
    FNS_atk = FNSack(config=config, neter=neter, dataer=dataer)

    ## for OnePixel attack method
    OnePixel_atk = OnePixel(neter.net, neter, pixels=1, steps=200, popsize=50)

    ## for SimBA attack method
    SimBA_atk = SimBA(neter=neter)

    ## for FGSM attack method
    fgsm_atk = FGSM(model=neter.net, eps=config['epsilon'])

    ## for PGD attack method
    pgd_atk = PGD(model=neter.net, eps=config['epsilon'], alpha=config['epsilon'] / 5, steps=20)

    ########################################################################
    ## Experimet Type
    ########################################################################

    if config['exp_type'] == 'untarget':
        
        exper = ExpUntarget(config=config, dataer=dataer, neter=neter)
    
    elif config['exp_type'] == 'target':
        
        exper = ExpTarget(config=config, dataer=dataer, neter=neter)
    
    elif config['exp_type'] == 'param':

        pass
    
    elif config['exp_type'] == 'robust':

        pass

    elif config['exp_type'] == 'transfer':

        exper = ExpTransfer(config=config, dataer=dataer, neter=neter)

    elif config['exp_type'] == 'multiple':

        exper = ExpMUAPs(config, dataer, neter)
    
    elif config['exp_type'] == 'highdim':

        exper = ExpHighDim(config, dataer, neter)

    elif config['exp_type'] == 'gaussian':

        exper = ExpUntarget(config=config, dataer=dataer, neter=neter)

    else:

        pass
    

    ########################################################################
    ## Experimet Case & Attack Objection
    ########################################################################

    attack_obj = config['exp_obj']

    if config['exp_case'] == 'case_1_2':

        if attack_obj == 'FGSM':

            exper.Case_1_2(seed=0, way='FGSM', atk=fgsm_atk, default_num=config['attack_num'])

        elif attack_obj == 'OnePixel':

            exper.Case_1_2(seed=0, way='OnePixel', atk=OnePixel_atk, default_num=config['attack_num'])

        elif attack_obj == 'FNS':

            exper.Case_1_2(seed=0, way='FNS', atk=FNS_atk, default_num=config['attack_num'])   

        elif attack_obj == 'SimBA':

            exper.Case_1_2(seed=0, way='SimBA', atk=SimBA_atk, default_num=config['attack_num'])

        elif attack_obj == 'MPack':
            
            exper.Case_1_2(seed=0, way='MPack', atk=atk, default_num=config['attack_num'])
        
        elif attack_obj == 'GA':

            exper.Case_1_2(seed=0, way='GA', atk=atk, default_num=config['attack_num'])

        elif attack_obj == 'SMPack':

            exper.Case_1_2(seed=0, way='SMPack', atk=atk, default_num=config['attack_num'])
        
        elif attack_obj == 'PGD':

            exper.Case_1_2(seed=0, way='PGD', atk=pgd_atk, default_num=config['attack_num'])
        
        elif attack_obj == 'export':

            exper.save_total_file(case='Case1')
            exper.save_total_file(case='Case2')

        else:

            raise Exception('No such attack method !')

        exper.show_result('Case1')
        exper.show_result('Case2')

    elif config['exp_case'] == 'case_3':

        exper.Case_3(atk=atk)

    elif config['exp_case'] == 'pair':
        
        exper.generate_attack_pair(atk=atk)

    elif config['exp_case'] == 'norm':

        exper.norm_calculate(atk=atk)

    elif config['exp_case'] == 'show':

        exper.show_result('Case1')
        exper.show_result('Case2')

    elif config['exp_case'] == 'latex':

        exper.show_result('Case1')

        exper.show_result('Case2')

        exper.latex_export(case='Case1')

        exper.latex_export(case='Case2')

    elif config['exp_case'] == 'matrix':

        exper.plot_matrix(atk=atk, way=attack_obj)


    elif config['exp_case'] == 'muaps':

        ploter = Plot()
        ploter.plot_UAPs(
            atk.get_untarget_MUAPs(
                source_class=3,
                epsilon=config['epsilon'],
                norm_type=config['attck_norm'],
                method=config['attck_method'],
                batch=config['base_batch'],
                size=config['base_size'],
                mode=config['base_mode']
            ),
        )
    
    elif config['exp_case'] == 'muaps_asr':

        exper.attck_acc(atk, 0)

    elif config['exp_case'] == 'performance':

        neter.Test_model()
        neter.Robust_test_model(fgsm_atk)
        neter.Noise_test_model()

    elif config['exp_case'] == 'case_transfer_untarget':

        exper.Case_transfer_MPack(seed=0, way='MPack', atk=atk, default_num=config['attack_num'])
        exper.show_result('Untarget')
    
    elif config['exp_case'] == 'case_transfer_target':
        exper.Case_transfer_target_MPack(seed=0, way='MPack', atk=atk, default_num=config['attack_num'])
        exper.show_result('Target')

    elif config['exp_case'] == 'case_transfer_speed':

        # exper.Case_speed(seed=0, way=config['exp_obj'], atk=atk)
        # exper.show_result('Speed')
        exper.Speed_plot()

    elif config['exp_case'] == 'case_transfer_target_speed':

        # exper.Case_target_speed(seed=0, way=config['exp_obj'], atk=atk)
        # exper.show_result('Speed-target')
        exper.Speed_target_plot()

    elif config['exp_case'] == 'case_highdim_untarget':

        if attack_obj == 'FGSM':

            exper.untarget(seed=0, way='FGSM', atk=fgsm_atk, default_num=config['attack_num'])

        elif attack_obj == 'OnePixel':

            exper.untarget(seed=0, way='OnePixel', atk=OnePixel_atk, default_num=config['attack_num'])

        elif attack_obj == 'FNS':

            exper.untarget(seed=0, way='FNS', atk=FNS_atk, default_num=config['attack_num'])   

        elif attack_obj == 'SimBA':

            exper.untarget(seed=0, way='SimBA', atk=SimBA_atk, default_num=config['attack_num'])

        elif attack_obj == 'MPack':
            
            exper.untarget(seed=0, way='MPack', atk=atk, default_num=config['attack_num'])
        
        elif attack_obj == 'GA':

            exper.untarget(seed=0, way='GA', atk=atk, default_num=config['attack_num'])

        elif attack_obj == 'SMPack':

            exper.untarget(seed=0, way='SMPack', atk=atk, default_num=config['attack_num'])
        
        elif attack_obj == 'PGD':

            exper.untarget(seed=0, way='PGD', atk=pgd_atk, default_num=config['attack_num'])
        
        elif attack_obj == 'pass':

            pass

        else:

            raise Exception('No such attack method !')

        exper.show_result('Untarget_Case1')
        exper.show_result('Untarget_Case2')

    elif config['exp_case'] == 'case_highdim_target':

        if attack_obj == 'FGSM':

            exper.target(seed=0, way='FGSM', atk=fgsm_atk, default_num=config['attack_num'])

        elif attack_obj == 'OnePixel':

            exper.target(seed=0, way='OnePixel', atk=OnePixel_atk, default_num=config['attack_num'])

        elif attack_obj == 'FNS':

            exper.target(seed=0, way='FNS', atk=FNS_atk, default_num=config['attack_num'])   

        elif attack_obj == 'SimBA':

            exper.target(seed=0, way='SimBA', atk=SimBA_atk, default_num=config['attack_num'])

        elif attack_obj == 'MPack':
            
            exper.target(seed=0, way='MPack', atk=atk, default_num=config['attack_num'])
        
        elif attack_obj == 'GA':

            exper.target(seed=0, way='GA', atk=atk, default_num=config['attack_num'])

        elif attack_obj == 'SMPack':

            exper.target(seed=0, way='SMPack', atk=atk, default_num=config['attack_num'])
        
        elif attack_obj == 'PGD':

            exper.target(seed=0, way='PGD', atk=pgd_atk, default_num=config['attack_num'])
        
        elif attack_obj == 'pass':

            pass

        else:

            raise Exception('No such attack method !')

        exper.show_result('Target_Case1')
        exper.show_result('Target_Case2')
    
    elif config['exp_case'] == 'high_dim_show':

        exper.show(atk=atk)

    elif config['exp_case'] == 'case_param_base_size':

        exper.base_size_plot()

    elif config['exp_case'] == 'case_param_base_batch':

        exper.base_batch_plot()

    elif config['exp_case'] == 'pass':

        pass

    else:

        pass



    ##### plot
    # ploter = Plot()
    # ploter.plot_UAPs(
    #     perturbations=atk.get_untarget_MUAPs(
    #         source_class=3,
    #         epsilon=config['epsilon'],
    #         norm_type=config['attck_norm'],
    #         method=config['attck_method'],
    #         batch=config['base_batch'],
    #         size=config['base_size'],
    #         mode=config['base_mode']
    #     ),
    #     epsilon=config['epsilon'],
    #     name='test',
    # )


#########################################################################
## Untarget Attck                                                      ##
#########################################################################
# source_class = config['source_class']

# # class_loader = dataer.get_class_loader(class_index=source_class, isTrain=False, batch_size=1)
# loader = dataer.get_loader(batch_size=1, isTrain=False, isShuffle=True)

# succ_number = 0

# for index, (image, label) in enumerate(loader):

#     print('Attacking {}-th sample'.format(index))
#     flag, adv_image = atk.untarget_attack(
#                         origin_sample=image,
#                         source_class=label[0],
#                         epsilon=config['epsilon'],
#                         norm_type=config['attck_norm'],
#                         method=config['attck_method'],
#                         batch=config['base_batch'],
#                         size=config['base_size'],
#                         mode=config['base_mode'],
#                     )
#     if flag:
#         succ_number += 1
#     if index == 99:
#         break

# print('Untarget Attack ASR(Randomly 100 samples without the target class): {}%'.format(succ_number))

#########################################################################
## Target Random-Target Attck                                          ##
#########################################################################
# target_class = config['target_class']

# exclude_class_loader = dataer.get_exclude_class_loader(class_index=target_class, isTrain=False, batch_size=1)

# succ_number = 0

# for index, (image, label) in enumerate(exclude_class_loader):

#     print('Attacking {}-th sample'.format(index))
#     flag, adv_image = atk.target_attck(
#                         origin_sample=image,
#                         source_class=label[0],
#                         target_class=target_class,
#                         epsilon=config['epsilon'],
#                         norm_type=config['attck_norm'],
#                         method=config['attck_method'],
#                         batch=config['base_batch'],
#                         size=config['base_size'],
#                         mode=config['base_mode'],
#                     )
#     if flag:
#         succ_number += 1
#     if index == 99:
#         break

# print('Target Attack ASR(Randomly 100 samples without the target class): {}%'.format(succ_number))


#########################################################################
## Target Source Target Attck                                          ##
#########################################################################

# source_class = config['source_class']
# target_class = config['target_class']

# class_loader = dataer.get_class_loader(class_index=source_class, isTrain=False, batch_size=1)


# count = 0
# total = 0
# for image, label in class_loader:

#     print('Attck {}-th sample'.format(total))

#     flag, adv_image = atk.target_attck(
#                         origin_sample=image,
#                         source_class=source_class,
#                         target_class=target_class,
#                         epsilon=config['epsilon'],
#                         norm_type=config['attck_norm'],
#                         method=config['attck_method'],
#                         batch=config['base_batch'],
#                         size=config['base_size'],
#                         mode=config['base_mode'],   
#                     )
#     total += 1
#     if total == 1:
#         break
#     if flag:
#         count += 1

# print('ASR : {:.2f}%'.format(count / total * 100))


